/*
 *  Copyright (c) 2012 The WebRTC project authors. All Rights Reserved.
 *
 *  Use of this source code is governed by a BSD-style license
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
 *  in the file PATENTS.  All contributing project authors may
 *  be found in the AUTHORS file in the root of the source tree.
 */

#include <assert.h>
#include <math.h>
#include <string.h>
#include <stdlib.h>

#include "fft4g.h"
#include "include/dynamic_range_compression.h"
#include "drc_core.h"
#include "windows_private.h"

// Initialize state.
int Drc_InitCore(DynamicRangeCompressionC* self, uint32_t fs, int threshold1, float ratio, int threshold2, int mode, float supPoint, float supExp) {
  int i;
  // Check for valid pointer.
  if (self == NULL) {
    return -1;
  }

  if (threshold2 < 0 || ratio < 0 || ratio > 1 || threshold2 > threshold1) {
    return -1;
  }
  
  // Initialization of struct.
  if (fs == 8000 || fs == 16000) {
    self->fs = fs;
  } else {
    return -1;
  }
  self->windShift = 0;
  
  // We only support 10ms frames.
  if (fs == 8000) {
    self->blockLen = 80;
    self->drcLen = 128;
    self->window = kBlocks80w128;
  } else {
    self->blockLen = 160;
    self->drcLen = 256;
    self->window = kBlocks160w256;
  }
  
  self->magnLen = self->drcLen / 2 + 1;  // Number of frequency bins.

  self->threshold1 = threshold1;
  self->ratio = ratio;
  self->threshold2 = threshold2;
  /* hard knee */
  self->mode = mode;
  
  self->supPoint = supPoint;
  self->supExp = supExp;

  // Initialize FFT work arrays.
  self->ip[0] = 0;  // Setting this triggers initialization.
  memset(self->dataBuf, 0, sizeof(float) * DRC_BLOCKL_MAX);
  Drc_rdft(self->drcLen, 1, self->dataBuf, self->ip, self->wfft);

  memset(self->dataBuf, 0, sizeof(float) * DRC_BLOCKL_MAX);
  memset(self->syntBuf, 0, sizeof(float) * DRC_BLOCKL_MAX);

  /* self->updates = 0; */

  self->initFlag = 1;
  return 0;
}

// Updates |buffer| with a new |frame|.
// Inputs:
//   * |frame| is a new speech frame or NULL for setting to zero.
//   * |frame_length| is the length of the new frame.
//   * |buffer_length| is the length of the buffer.
// Output:
//   * |buffer| is the updated buffer.
static void UpdateBuffer(const float* frame,
                         size_t frame_length,
                         size_t buffer_length,
                         float* buffer) {
  assert(buffer_length < 2 * frame_length);

  memcpy(buffer,
         buffer + frame_length,
         sizeof(*buffer) * (buffer_length - frame_length));
  if (frame) {
    memcpy(buffer + buffer_length - frame_length,
           frame,
           sizeof(*buffer) * frame_length);
  } else {
    memset(buffer + buffer_length - frame_length,
           0,
           sizeof(*buffer) * frame_length);
  }
}

// Transforms the signal from time to frequency domain.
// Inputs:
//   * |time_data| is the signal in the time domain.
//   * |time_data_length| is the length of the analysis buffer.
//   * |magnitude_length| is the length of the spectrum magnitude, which equals
//     the length of both |real| and |imag| (time_data_length / 2 + 1).
// Outputs:
//   * |time_data| is the signal in the frequency domain.
//   * |real| is the real part of the frequency domain.
//   * |imag| is the imaginary part of the frequency domain.
//   * |magn| is the calculated signal magnitude in the frequency domain.
static void FFT(DynamicRangeCompressionC* self,
                float* time_data,
                size_t time_data_length,
                size_t magnitude_length,
                float* real,
                float* imag,
                float* magn) {
  size_t i;

  assert(magnitude_length == time_data_length / 2 + 1);

  Drc_rdft(time_data_length, 1, time_data, self->ip, self->wfft);

  imag[0] = 0;
  real[0] = time_data[0];
  magn[0] = fabsf(real[0]) + 1.f;
  imag[magnitude_length - 1] = 0;
  real[magnitude_length - 1] = time_data[1];
  magn[magnitude_length - 1] = fabsf(real[magnitude_length - 1]) + 1.f;
  for (i = 1; i < magnitude_length - 1; ++i) {
    real[i] = time_data[2 * i];
    imag[i] = time_data[2 * i + 1];
    // Magnitude spectrum.
    magn[i] = sqrtf(real[i] * real[i] + imag[i] * imag[i]) + 1.f;
  }
}

// Transforms the signal from frequency to time domain.
// Inputs:
//   * |real| is the real part of the frequency domain.
//   * |imag| is the imaginary part of the frequency domain.
//   * |magnitude_length| is the length of the spectrum magnitude, which equals
//     the length of both |real| and |imag|.
//   * |time_data_length| is the length of the analysis buffer
//     (2 * (magnitude_length - 1)).
// Output:
//   * |time_data| is the signal in the time domain.
static void IFFT(DynamicRangeCompressionC* self,
                 const float* real,
                 const float* imag,
                 size_t magnitude_length,
                 size_t time_data_length,
                 float* time_data) {
  size_t i;

  assert(time_data_length == 2 * (magnitude_length - 1));

  time_data[0] = real[0];
  time_data[1] = real[magnitude_length - 1];
  for (i = 1; i < magnitude_length - 1; ++i) {
    time_data[2 * i] = real[i];
    time_data[2 * i + 1] = imag[i];
  }
  Drc_rdft(time_data_length, -1, time_data, self->ip, self->wfft);

  for (i = 0; i < time_data_length; ++i) {
    time_data[i] *= 2.f / time_data_length;  // FFT scaling.
  }
}

// Windows a buffer.
// Inputs:
//   * |window| is the window by which to multiply.
//   * |data| is the data without windowing.
//   * |length| is the length of the window and data.
// Output:
//   * |data_windowed| is the windowed data.
static void Windowing(const float* window,
                      const float* data,
                      size_t length,
                      float* data_windowed) {
  size_t i;

  for (i = 0; i < length; ++i) {
    data_windowed[i] = window[i] * data[i];
  }
}

void Drc_ProcessCore(DynamicRangeCompressionC* self,
                          const float* const* speechFrame,
                          float* const* outFrame) {
  // Main routine for noise reduction.
  size_t i, j;

  float fout[BLOCKL_MAX];
  float winData[DRC_BLOCKL_MAX];
  float magn[HALF_DRC_BLOCKL];
  float real[DRC_BLOCKL_MAX], imag[HALF_DRC_BLOCKL];

  // Check that initiation has been done.
  assert(self->initFlag == 1);

  // Update analysis buffer for L band.
  UpdateBuffer(speechFrame[0], self->blockLen, self->drcLen, self->dataBuf);

  Windowing(self->window, self->dataBuf, self->drcLen, winData);

  FFT(self, winData, self->drcLen, self->magnLen, real, imag, magn);
  /* printf("self->magnLen=%d\n", self->magnLen); */
  int supPoint = self->magnLen;
  if(self->supPoint < 1)
      supPoint = (int)(self->magnLen*self->supPoint);

  for (i = 0; i < self->magnLen; i++) {
    // Flooring bottom.
    double E = sqrt(real[i]*real[i] + imag[i]*imag[i]);
    double scale = 1.0;
    if(E > self->threshold1 && self->mode & 1){
	scale = ((E-self->threshold1)*self->ratio+self->threshold1)/E;
    }else if (E < self->threshold2 && self->mode & 2){
	scale = (E/self->threshold2);
    }
    if(i>supPoint){
	float a = (float)(self->magnLen - i - 1)/(float)(self->magnLen - supPoint - 1);
    	scale *= pow(a, self->supExp);;
    }
    if(self->mode & 4){
	if(E <= self->threshold2)
	{
	    scale = scale*0.0;
	}
	else
	{
	    scale = scale*(E-self->threshold2)/E;
	}
    }
    /* if(i<3) scale = 0; */
    real[i] *= scale;
    imag[i] *= scale;
  }
  
  IFFT(self, real, imag, self->magnLen, self->drcLen, winData);

  Windowing(self->window, winData, self->drcLen, winData);

  // Synthesis.
  for (i = 0; i < self->drcLen; i++) {
    self->syntBuf[i] += winData[i];
  }
  // Read out fully processed segment.
  for (i = self->windShift; i < self->blockLen + self->windShift; i++) {
    fout[i - self->windShift] = self->syntBuf[i];
  }
  // Update synthesis buffer.
  UpdateBuffer(NULL, self->blockLen, self->drcLen, self->syntBuf);

  for (i = 0; i < self->blockLen; ++i)
    outFrame[0][i] =
        DRC_SAT(DRC_WORD16_MAX, fout[i], DRC_WORD16_MIN);

}
